package study;
import java.util.Comparator;
import java.util.PriorityQueue;

public class code04_TopKSumCrossTwoArrays {

    // 放入大根堆里的结构
    public static class Node {
        public int index1;// arr1中的位置
        public int index2;// arr2中的位置
        public int sum; // arr1[index1] + arr2[index2]

        public Node(int i1,int i2,int s){
            index1 = i1;
            index2 = i2;
            sum = s;
        }
    }

    // 大根堆的比较器
    public static class  MaxHeapCom implements Comparator<Node> {
        @Override
        public int compare(code04_TopKSumCrossTwoArrays.Node o1, code04_TopKSumCrossTwoArrays.Node o2) {
            return o2.sum - o1.sum;
        }
    }

    public static int[] topKSum(int[] arr1,int[] arr2,int topk){
        if (arr1 == null || arr2 == null || topk<1) {
            return null;
        }
        // 防止topk出现捣乱数据，arr组合 数据有限
        topk = Math.min(topk, arr1.length*arr2.length);

        int[] res = new int[topk];
        int resIndex = 0;

        PriorityQueue<Node> maxHeap = new PriorityQueue<>(new MaxHeapCom());

        // set[i][j] == false, 代表 arr1[i] arr2[j] 之前没使用过
        // set[i][j] == true, 代表 arr1[i] arr2[j] 之前被使用
        boolean[][] set = new boolean[arr1.length][arr2.length];
        int i1 = arr1.length;
        int i2 = arr2.length;
        maxHeap.add(new Node(i1, i2, arr1[i1] + arr2[i2]));
        set[i1][i2] = true;
        while (resIndex != topk) {
            Node curNode = maxHeap.poll();
            res[resIndex++] = curNode.sum;
            i1 = curNode.index1;
            i2 = curNode.index2;
            if (i1 - 1 >= 0 && !set[i1 - 1][i2]) {
                set[i1 - 1][i2] = true;
                maxHeap.add(new Node(i1 - 1, i2, arr1[i1-1] + arr2[i2]));
            }
             if (i2 - 1 >= 0 && !set[i1][i2-1]) {
                set[i1][i2 - 1] = true;
                maxHeap.add(new Node(i1, i2 - 1, arr1[i1] + arr2[i2 - 1]));
            }
            
        }
        return res;
    }
}
